package cn.kinoko.common.aspect;

import cn.kinoko.common.constant.HeaderConstant;
import cn.kinoko.common.utils.ExceptionUtil;
import cn.kinoko.common.utils.SpelUtil;
import cn.kinoko.common.utils.ThreadLocalUtil;
import cn.kinoko.common.utils.TokenBucketLimiter;
import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.RemovalCause;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.stereotype.Component;

import java.time.Duration;
import java.util.concurrent.ExecutorService;

/**
 * 分布式缓存切面
 *
 * @author kk
 */
@Aspect
@Slf4j
@Component
public class TokenBucketLimiterAspect implements InitializingBean {

    @Resource
    private ExecutorService executorService;

    private Cache<String, TokenBucketLimiter> limiterCache;

    @Override
    public void afterPropertiesSet() {
        this.limiterCache = Caffeine.newBuilder()
                // 5分钟未使用则过期
                .expireAfterAccess(Duration.ofSeconds(5))
                .removalListener(TokenBucketLimiterAspect::onRemoval)
                .executor(executorService)
                .build();
    }

    /**
     * 缓存失效回调，关闭线程池，防止线程池泄露
     *
     * @param key   key
     * @param value value
     * @param cause 原因
     */
    private static void onRemoval(@Nullable String key, @Nullable TokenBucketLimiter value, RemovalCause cause) {
        if (value != null) {
            value.shutdown();
            log.info("[caffeine]TokenBucketLimiter: {} 缓存失效，回收线程池", key);
        }
    }

    @Around("@annotation(limiter)")
    public Object around(ProceedingJoinPoint joinPoint, cn.kinoko.common.aspect.annotation.TokenBucketLimiter limiter) throws Throwable {
        // 获取redisKey
        String redisKey = SpelUtil.getRedisKey(joinPoint, limiter.key(), limiter.params());
        if (limiter.currentUser()) {
            redisKey += (":" + ThreadLocalUtil.get(HeaderConstant.UID_HEADER));
        }
        // 获取限流器，没有则构建，并且保证只有一个线程参与构建
        TokenBucketLimiter tokenBucketLimiter = limiterCache.get(redisKey, k -> new TokenBucketLimiter(limiter.capacity(), limiter.rate(), limiter.waitTime()));
        // 尝试获取令牌
        ExceptionUtil.assertion(!tokenBucketLimiter.tryConsume(), limiter.errorDesc());
        // 执行业务方法
        return joinPoint.proceed();
    }

}
